package llm

import (
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

/* 百度千帆 */

const (
	qianfan_app_id     = ""
	qianfan_api_key    = ""
	qianfan_secret_key = ""

	baidu_token_url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=" +
		qianfan_api_key + "&client_secret=" + qianfan_secret_key
)

var ernie_urls = map[string]string{
	MODEL_ERNIE_BOT:   "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token=",
	MODEL_ERNIE_BOT_4: "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=",
}

/* ERNIE-Bot */
type ErnieBot struct {
	Name string
	*Token
}

func (e *ErnieBot) RefreshToken() error {
	if e.expire > time.Now().Unix() {
		return nil
	}

	rsp, err := client.R().Get(baidu_token_url)
	if err != nil {
		return err
	}

	if !rsp.IsSuccess() {
		return errors.New(string(rsp.Body()))
	}

	var v BaiduAccessToken
	err = json.Unmarshal(rsp.Body(), &v)
	if err != nil {
		return err
	}

	e.token = v.AccessToken
	e.expire = time.Now().Unix() + v.ExpiresIn - 5

	return nil
}

func (e *ErnieBot) CallApi(system_prompt, user_prompt string) (string, error) {

	result, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetBody(generateErnieReq(system_prompt, user_prompt)).
		Post(ernie_urls[e.Name] + e.token)

	if err != nil {
		return "", err
	}

	var rsp ErnieBotRsp
	err = json.Unmarshal(result.Body(), &rsp)
	if err != nil {
		return "", err
	}

	if rsp.Flag != 0 {
		return "", fmt.Errorf("%s invoke failure", e.Name)
	}

	return rsp.Result, nil

}

func generateErnieReq(system_prompt, user_prompt string) map[string]interface{} {
	return map[string]interface{}{
		"system": system_prompt,
		"messages": []map[string]string{
			{
				"role":    "user",
				"content": user_prompt,
			},
		},
		"response_format": "json_object",
		"disable_search":  true,
	}
}

type BaiduAccessToken struct {
	AccessToken string `json:"access_token"`
	ExpiresIn   int64  `json:"expires_in"`
}

type ErnieUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type ErnieBotRsp struct {
	ID         string     `json:"id"`
	Created    int        `json:"created"`
	SentenceID int        `json:"sentence_id"`
	Result     string     `json:"result"`
	Flag       int        `json:"flag"`
	Usage      ErnieUsage `json:"usage"`
}
